import numpy as np
import pickle
from scipy.spatial.transform import Rotation as R
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

def vis_trans(trans):

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    ax.scatter(trans[:, 0], trans[:, 1], trans[:, 2])

    ax.set_xlabel('X Label')
    ax.set_ylabel('Y Label')
    ax.set_zlabel('Z Label')

    plt.show()

def vis_rot(rots):
    fig = plt.figure(figsize=(8, 8))
    ax = fig.add_subplot(111, projection='3d')
    for rot in rots:
        rotated_vector  = np.dot(rot, np.array([0, 0, 1]))
        x, y, z = rotated_vector / np.linalg.norm(rotated_vector)
        ax.scatter(x, y, z, color='b')
    # Set plot limits to make the sphere look like a sphere
    ax.set_xlim([-1, 1])
    ax.set_ylim([-1, 1])
    ax.set_zlim([-1, 1])

    # Set equal scaling for all axes
    ax.set_box_aspect([1,1,1])
    plt.title('Visualization of Quaternion Rotation Data Points on a Sphere')
    plt.show()

def tf_to_xyzquat_numpy(pose):
    """
    convert 4 x 4 transformation matrices to [x, y, z, qx, qy, qz, qw]
    """
    pose = np.atleast_3d(pose)

    r = R.from_matrix(np.array(pose[:, 0:3, 0:3]))
    q = r.as_quat()  # qx, qy, qz, qw
    t = pose[:, :3, 3]
    xyz_quat = np.concatenate((t, q), axis=1)

    return xyz_quat  # (N, 7)

data_buffer = []

with open ('obj_2.urdf.pkl', 'rb') as f:
    data = pickle.load(f)

gripper_width = data['gripper_width'] # List[[left_finger_joint_dof, right_finger_joint_dof],...]
fingertip_center_pose = data['fingertip_center_pose'] # List[[x, y, z, qx, qy, qz, qw],...] the panda_fingertip_centered link defined in urdf
obj_pose = data['obj_pose'] # List[[x, y, z, qx, qy, qz, qw],...] the object pose in local env frame
left_finger_pose = data['left_finger_pose'] # List[[x, y, z, qx, qy, qz, qw],...] the panda_leftfinger link defined in urdf
right_finger_pose = data['right_finger_pose'] # List[[x, y, z, qx, qy, qz, qw],...] the panda_rightfinger link defined in urdf
obj2gripper = data['obj2gripper'] # List[[4x4 tf matrix],...] the object pose relative to fingertip_center_pose (obj2fingertip_centered)
franka_joint_pose = data['franka_joint_pose'] # the joint pose of franka arm
obj_name = data['obj_name']

# vis the trans/rot distribution of obj2gripper
obj2gripper_list = data['obj2gripper']
for obj2gripper in obj2gripper_list:
    data_buffer.extend(obj2gripper)
data_buffer = np.array(data_buffer, dtype=np.float32)

trans = data_buffer[:, 0:3, 3]
rots = data_buffer[:, 0:3, 0:3]
vis_trans(trans=trans)
vis_rot(rots=rots)
